#!/bin/bash

# Milvus Docker Standalone Management Script
# Usage: ./manage.sh [command]

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check if Docker is running
check_docker() {
    if ! docker info >/dev/null 2>&1; then
        log_error "Docker未运行，请先启动Docker服务"
        exit 1
    fi
}

# Check if Docker Compose is available
check_docker_compose() {
    if ! command -v docker-compose >/dev/null 2>&1; then
        log_error "Docker Compose未安装"
        exit 1
    fi
}

# Start Milvus services
start_milvus() {
    log_info "启动Milvus服务..."
    
    # Create volumes directory if it doesn't exist
    mkdir -p volumes/{milvus,etcd,minio}
    
    docker-compose up -d
    
    log_info "等待服务启动..."
    sleep 10
    
    # Wait for Milvus to be healthy
    local max_attempts=30
    local attempt=1
    
    while [ $attempt -le $max_attempts ]; do
        if curl -f http://localhost:9091/healthz >/dev/null 2>&1; then
            log_success "Milvus服务已就绪！"
            break
        fi
        
        log_info "尝试 $attempt/$max_attempts: 等待Milvus服务启动..."
        sleep 5
        ((attempt++))
    done
    
    if [ $attempt -gt $max_attempts ]; then
        log_error "Milvus服务启动超时"
        show_logs
        exit 1
    fi
    
    show_status
}

# Stop Milvus services
stop_milvus() {
    log_info "停止Milvus服务..."
    docker-compose down
    log_success "Milvus服务已停止"
}

# Restart Milvus services
restart_milvus() {
    log_info "重启Milvus服务..."
    stop_milvus
    start_milvus
}

# Show service status
show_status() {
    log_info "服务状态:"
    docker-compose ps
    
    echo
    log_info "服务访问地址:"
    echo "  Milvus服务: localhost:19530"
    echo "  Milvus健康检查: http://localhost:9091/healthz"
    echo "  MinIO控制台: http://localhost:9001 (minioadmin/minioadmin)"
}

# Show logs
show_logs() {
    if [ -n "$1" ]; then
        log_info "显示 $1 服务日志..."
        docker-compose logs -f "$1"
    else
        log_info "显示所有服务日志..."
        docker-compose logs -f
    fi
}

# Test Milvus connection
test_connection() {
    log_info "测试Milvus连接..."
    
    # Test health endpoint
    if curl -f http://localhost:9091/healthz >/dev/null 2>&1; then
        log_success "健康检查通过"
    else
        log_error "健康检查失败"
        return 1
    fi
    
    # Test with Python if available
    if command -v python3 >/dev/null 2>&1; then
        python3 -c "
try:
    from pymilvus import connections, utility
    connections.connect('default', host='localhost', port='19530')
    version = utility.get_server_version()
    print(f'✓ 已连接到Milvus版本: {version}')
except ImportError:
    print('⚠ pymilvus未安装。安装命令: pip install pymilvus')
except Exception as e:
    print(f'✗ 连接失败: {e}')
    exit(1)
"
    else
        log_warning "Python3不可用，无法进行高级测试"
    fi
}

# Backup data
backup_data() {
    log_info "备份Milvus数据..."
    
    if [ ! -d "volumes" ]; then
        log_error "未找到数据目录"
        exit 1
    fi
    
    local backup_name="milvus-backup-$(date +%Y%m%d-%H%M%S).tar.gz"
    tar -czf "$backup_name" volumes/
    
    log_success "备份完成: $backup_name"
}

# Clean up everything
cleanup() {
    log_warning "这将删除所有Milvus数据。确定要继续吗？(y/N)"
    read -r response
    if [[ "$response" =~ ^[Yy]$ ]]; then
        log_info "清理Milvus部署..."
        docker-compose down -v
        rm -rf volumes/
        log_success "清理完成"
    else
        log_info "清理已取消"
    fi
}

# Show help
show_help() {
    echo "Milvus Docker Standalone 管理脚本"
    echo
    echo "用法: $0 [命令]"
    echo
    echo "命令:"
    echo "  start     启动Milvus服务"
    echo "  stop      停止Milvus服务"
    echo "  restart   重启Milvus服务"
    echo "  status    显示服务状态"
    echo "  logs      显示所有服务日志"
    echo "  logs <service>  显示指定服务日志"
    echo "  test      测试Milvus连接"
    echo "  backup    备份数据"
    echo "  cleanup   清理所有数据和容器"
    echo "  help      显示帮助信息"
    echo
    echo "示例:"
    echo "  $0 start"
    echo "  $0 logs milvus"
    echo "  $0 test"
}

# Main script logic
main() {
    check_docker
    check_docker_compose
    
    case "${1:-help}" in
        start)
            start_milvus
            ;;
        stop)
            stop_milvus
            ;;
        restart)
            restart_milvus
            ;;
        status)
            show_status
            ;;
        logs)
            show_logs "$2"
            ;;
        test)
            test_connection
            ;;
        backup)
            backup_data
            ;;
        cleanup)
            cleanup
            ;;
        help|--help|-h)
            show_help
            ;;
        *)
            log_error "未知命令: $1"
            show_help
            exit 1
            ;;
    esac
}

# Run main function with all arguments
main "$@"
